"""数据加载和标签处理模块"""
import os
import glob
import numpy as np
from typing import List, Tuple, Optional

class LabelLoader:
    """标签加载器"""
    
    @staticmethod
    def load_label(label_path: str) -> np.ndarray:
        """
        加载YOLO格式标签文件
        
        Args:
            label_path: 标签文件路径
            
        Returns:
            标签数组 [class_id, x_center, y_center, width, height]
        """
        if not os.path.exists(label_path) or os.path.getsize(label_path) == 0:
            return np.zeros((0, 5))
        
        try:
            labels = np.loadtxt(label_path).reshape(-1, 5)
            return labels
        except Exception as e:
            print(f"Warning: Failed to load label {label_path}: {e}")
            return np.zeros((0, 5))

class DatasetLoader:
    """数据集加载器"""
    
    def __init__(self, img_dir: str, label_dir: str, img_ext: str = '*.jpg'):
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.img_ext = img_ext
        self._validate_dirs()
    
    def _validate_dirs(self):
        """验证目录是否存在"""
        if not os.path.exists(self.img_dir):
            raise FileNotFoundError(f"Image directory not found: {self.img_dir}")
        if not os.path.exists(self.label_dir):
            raise FileNotFoundError(f"Label directory not found: {self.label_dir}")
    
    def get_image_paths(self) -> List[str]:
        """获取所有图像路径"""
        pattern = os.path.join(self.img_dir, self.img_ext)
        return sorted(glob.glob(pattern))
    
    def get_label_path(self, img_path: str) -> str:
        """根据图像路径获取对应标签路径"""
        img_name = os.path.basename(img_path)
        label_name = os.path.splitext(img_name)[0] + '.txt'
        return os.path.join(self.label_dir, label_name)
    
    def load_image_label_pair(self, img_path: str) -> Tuple[str, np.ndarray]:
        """加载图像标签对"""
        label_path = self.get_label_path(img_path)
        labels = LabelLoader.load_label(label_path)
        return img_path, labels